package cn.kduck;

import io.modelcontextprotocol.client.McpSyncClient;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.memory.InMemoryChatMemory;
import org.springframework.ai.chat.messages.AssistantMessage.ToolCall;
import org.springframework.ai.chat.model.Generation;
import org.springframework.ai.mcp.SyncMcpToolCallbackProvider;
import org.springframework.http.MediaType;
import org.springframework.util.StringUtils;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.text.SimpleDateFormat;
import java.util.Date;
import java.util.List;
import java.util.UUID;

@RestController
@RequestMapping("/sse")
public class SseController {


    private ChatClient chatClient;
    private SseEmitter emitter;

    public SseController(ChatClient.Builder builder, List<McpSyncClient> mcpClients) {
        chatClient = builder
        .defaultAdvisors(new MessageChatMemoryAdvisor(new InMemoryChatMemory()))
        //这里装配可用的工具
        .defaultTools(new SyncMcpToolCallbackProvider(mcpClients))
        .build();
    }

    @GetMapping(path = "/events", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public SseEmitter handleSse() {
        emitter = new SseEmitter(0L);

        sendMessage(emitter,"🤖️助手已准备好了，可以开始对话了\n");

        emitter.onCompletion(() -> System.out.println("SSE completed"));
        emitter.onTimeout(() -> System.out.println("SSE timeout"));
        emitter.onError((ex) -> System.out.println("SSE error: " + ex.getMessage()));

        return emitter;
    }

    @PostMapping(path = "/send")
    public void send(@RequestParam("message") String message){

        sendMessage(emitter,"\n\n👨你的问题：<span style=\"color:blue\">"+message+"</span>\n");
        sendMessage(emitter,"🤖️助手的回答：\n");
//        chatClient.prompt(message)
//                .stream()
//                .content()
//                .subscribe(response -> {
////                    System.out.print(response);
//                    try {
//                        emitter.send(SseEmitter.event()
//                                .id(String.valueOf(UUID.randomUUID().toString()))
//                                .name("message")
//                                .data(response));
//                    } catch (IOException e) {
//                        emitter.completeWithError(e);
//                        throw new RuntimeException(e);
//                    }
//                });

        chatClient.prompt(message)
                .stream()
                .chatResponse().subscribe(response -> {
                    try {
                        Generation result = response.getResult();
                        emitter.send(SseEmitter.event()
                                .id(String.valueOf(UUID.randomUUID().toString()))
                                .name("message")
                                .data(result.getOutput().getText()));
//                        if(response.hasToolCalls()){
//                            System.out.println("🔧工具调用：" + response.getResult().getOutput().getToolCalls().stream().map(ToolCall::name));
//                        }
                        String finishReason = result.getMetadata().getFinishReason();
                        if(StringUtils.hasText(finishReason)) {
                            SimpleDateFormat sdf = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss");
                            System.out.println(sdf.format(new Date()) + "🆗输出停止：" + finishReason);
                        }
                    } catch (IOException e) {
                        emitter.completeWithError(e);
                        throw new RuntimeException(e);
                    }
                });
    }

    private void sendMessage(SseEmitter emitter,String object) {
        try {
            emitter.send(SseEmitter.event()
                    .id(String.valueOf(UUID.randomUUID().toString()))
                    .name("message")
                    .data(object));
        } catch (IOException e) {
            emitter.completeWithError(e);
            throw new RuntimeException(e);
        }
    }
}